from langchain.chains import RetrievalQA
from langchain.callbacks.streaming_stdout import StreamingStdOutCallbackHandler
from langchain.callbacks.manager import CallbackManager
from langchain_community.llms import Ollama
from langchain_community.embeddings.ollama import OllamaEmbeddings
from langchain_community.vectorstores import Chroma
from langchain.text_splitter import RecursiveCharacterTextSplitter
from langchain_community.document_loaders import DirectoryLoader
from langchain_community.document_loaders import TextLoader
from langchain.prompts import PromptTemplate
from langchain.memory import ConversationBufferMemory
import os
import time

with open("modelname.cfg","r") as f:
    modelname=f.readline().strip()
print("using llm modle",modelname)

#loader = DirectoryLoader('files', glob="**/*.txt")
loader = TextLoader("files/kcjlczyheadtime.md",autodetect_encoding=True)
data = loader.load()

print("pfdload done len of data is",len(data))

# Initialize text splitter
text_splitter = RecursiveCharacterTextSplitter(
    chunk_size=400,
    chunk_overlap=0,
    length_function=len
)
all_splits = text_splitter.split_documents(data)
print("text splited,got splits",len(all_splits))

n=0
run=False
for sptext in all_splits:
    print(sptext)
    n+=1
    if(n%500==0 and run !=True):
        print("press any key to continue,c for run",n,len(all_splits))
        x=input()
        if x=="c" or x=="C":
            run=True


# Create and persist the vector store
print("embedding using ollama model",modelname)
print("setup vectorstor")

vectorstore = Chroma.from_documents(persist_directory='jj',
    documents=all_splits,
    embedding=OllamaEmbeddings(model=modelname)
)
print("embeddong done")
vectorstore.persist()
print("vector store done")

docs=vectorstore.similarity_search("陈知远2024年04月学习了什么内容",k=5)

print(docs[0].page_content)

template = """You are a teacher for STEM course, here to help with questions of the user. Your tone should be professional and confident. Please addup teaching content by year,quater,and month.
    The knowlege you taugth is in the following technical fields:3D modeling,123D design,Python,minecraft MCPI,arduino,raspberry pi,C++,Java,robot,servo,html,serial port,risc-v,design thinking,innovation,project based learning PBL, teamwork,git,github,vscode,3D printing,laser cutting,
    CNC engraving,blockly,mixly,scratch,processing,app inventor,usaco,AP CS,patent,IP,robot arm,SLAM,ros2,uwb,tic-tac-toe,mediapipe,NLP,llm,ollama,RAG,langchain,CNN,pytorch,unitly,blender,linux,wsl,ubuntu,debian,device driver,soldering,
    circuit design,pcb layout,FPGA,ARM,openveno,home assistant,project management,sprint,agile,Dejango,matplotlib,numpy,opencv,pandas,chempy,chatgpt,openai,qwen,machine learning,RNN,pygame,etc.
    Context: {context}
    History: {history}
    User: {question}
    Chatbot:"""

llm = Ollama(model=modelname,verbose=True,callback_manager=CallbackManager([StreamingStdOutCallbackHandler()]))
